import torch,os
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection 
import pdb
from PIL import ImageDraw, ImageFont

DINO_CKP = '/home/shengjie/ckp/grounding-dino-base'
img_dir = '/mnt/nas/zhangshu/train_data/clothes_collection_clean/images'
yolo_save_dir = '/mnt/nas/shengjie/clothes_collection_clean/images_yolo'
os.makedirs(yolo_save_dir,exist_ok=True)
# img_file = '/data/shengjie/style_zhenzhi/img_1.jpg'
# img_file = os.path.join(
#     img_dir,os.listdir(img_dir)[0]
# )

model_id = DINO_CKP
device = "cuda" if torch.cuda.is_available() else "cpu"

text = "necklines.sleeve.pocket."
labels = text.split('.')[:-1] # 最后会有一个 '' 空字符串 要删掉
def get_label_id(label_text,labels):
    return labels.index(label_text)
# 转换函数
def xyxy_to_yolo(box, img_width, img_height):
    x_min, y_min, x_max, y_max = box
    x_center = (x_min + x_max) / 2 / img_width
    y_center = (y_min + y_max) / 2 / img_height
    width = (x_max - x_min) / img_width
    height = (y_max - y_min) / img_height
    return [x_center, y_center, width, height]


def visualize_results_pillow(image, results):
    draw = ImageDraw.Draw(image)
    
    # 加载字体（可选，需要系统支持）
    try:
        font = ImageFont.truetype("arial.ttf", 15)
    except:
        font = ImageFont.load_default()

    for box, label, score in zip(results['boxes'], results['labels'], results['scores']):
        x1, y1, x2, y2 = map(int, box.tolist())
        
        # 绘制边界框
        draw.rectangle([x1, y1, x2, y2], outline="green", width=2)
        
        # 添加标签和置信度
        text = f"{label}: {score:.2f}"
        draw.text((x1, y1 - 20), text, fill="red", font=font)
    
    return image

processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device)

from tqdm import tqdm
for img_name in tqdm(os.listdir(img_dir)):
    img_file = os.path.join(img_dir,img_name)

    assert os.path.exists(img_file)

    # image_url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    # image = Image.open(requests.get(image_url, stream=True).raw)
    image = Image.open(img_file)
    # Check for cats and remote controls
    # VERY important: text queries need to be lowercased + end with a dot


    inputs = processor(images=image, text=text, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)

    results = processor.post_process_grounded_object_detection(
        outputs,
        inputs.input_ids,
        box_threshold=0.2,
        text_threshold=0.4,
        target_sizes=[image.size[::-1]]
    )


    visualized_image = visualize_results_pillow(image.copy(), results[0])

    img_width, img_height = image.size
    # 生成 YOLO 格式标签
    yolo_lines = []
    for result in results:
        for score, label, box in zip(result['scores'], result['labels'], result['boxes']):
            # if score >= confidence_threshold and label in class_list:
            class_id = get_label_id(label,labels) if label in labels else 0
            yolo_box = xyxy_to_yolo(box.cpu().numpy(), img_width, img_height)
            line = f"{class_id} {yolo_box[0]:.6f} {yolo_box[1]:.6f} {yolo_box[2]:.6f} {yolo_box[3]:.6f}"
            yolo_lines.append(line)

    # 输出结果
    # print("\n".join(yolo_lines))

    yolo_save_path = os.path.join(yolo_save_dir,
                                os.path.splitext(img_name)[0]+'.txt') 
    
    # pdb.set_trace()
    with open(yolo_save_path,'w',encoding='utf-8') as f:
        f.write("\n".join(yolo_lines))

    visualized_image.save('tmp.jpg')
    # pdb.set_trace()
